import streamlit as st
import mysql.connector
from mysql.connector import Error
import pandas as pd
from datetime import datetime, date, time, timedelta
import hashlib


class FlightBookingSystem:
    def __init__(self, host, database, user, password):
        """初始化数据库连接"""
        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.connection = None

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                database=self.database,
                user=self.user,
                password=self.password,
                auth_plugin='mysql_native_password'
            )
            if self.connection.is_connected():
                print("成功连接到MySQL数据库")
                return True
        except Error as e:
            st.error(f"数据库连接失败: {e}")
            return False

    def disconnect(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            print("数据库连接已关闭")

    # 用户相关操作
    def add_user(self, username, password, email=None, phone=None,
                 first_name=None, last_name=None, date_of_birth=None, create_by=None):
        """添加新用户"""
        try:
            cursor = self.connection.cursor()
            password_hash = hashlib.sha256(password.encode()).hexdigest()
            query = """
            INSERT INTO users (username, password_hash, email, phone, 
                              first_name, last_name, date_of_birth, create_by)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """
            data = (username, password_hash, email, phone,
                    first_name, last_name, date_of_birth, create_by)
            cursor.execute(query, data)
            self.connection.commit()
            return cursor.lastrowid
        except Error as e:
            st.error(f"添加用户时出错: {e}")
            return None

    def get_user_by_id(self, user_id):
        """按ID查询用户"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM users WHERE user_id = %s", (user_id,))
            return cursor.fetchone()
        except Error as e:
            st.error(f"查询用户时出错: {e}")
            return None

    def get_all_users(self):
        """查询所有用户"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM users")
            return cursor.fetchall()
        except Error as e:
            st.error(f"查询用户时出错: {e}")
            return []

    def update_user(self, user_id, **kwargs):
        """更新用户信息"""
        try:
            cursor = self.connection.cursor()
            set_fields = []
            params = []

            if 'password' in kwargs:
                kwargs['password_hash'] = hashlib.sha256(kwargs.pop('password').encode()).hexdigest()

            for field, value in kwargs.items():
                if value is not None:
                    set_fields.append(f"{field} = %s")
                    params.append(value)

            if not set_fields:
                return False

            params.append(user_id)
            query = f"UPDATE users SET {', '.join(set_fields)} WHERE user_id = %s"
            cursor.execute(query, params)
            self.connection.commit()
            return cursor.rowcount > 0
        except Error as e:
            st.error(f"更新用户时出错: {e}")
            return False

    def delete_user(self, user_id):
        """删除用户"""
        try:
            cursor = self.connection.cursor()
            cursor.execute("DELETE FROM users WHERE user_id = %s", (user_id,))
            self.connection.commit()
            return cursor.rowcount > 0
        except Error as e:
            st.error(f"删除用户时出错: {e}")
            return False

    # 航班相关操作
    def add_flight(self, flight_number, departure_city, arrival_city, departure_time,
                   arrival_time, price, seats_available, create_by=None):
        try:
            cursor = self.connection.cursor()
            query = """
            INSERT INTO flights (flight_number, departure_city, arrival_city, 
                                departure_time, arrival_time, price, seats_available, create_by)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """
            data = (flight_number, departure_city, arrival_city, departure_time,
                    arrival_time, price, seats_available, create_by)
            cursor.execute(query, data)
            self.connection.commit()
            return cursor.lastrowid
        except Error as e:
            st.error(f"添加航班时出错: {e}")
            return None

    def get_flight_by_id(self, flight_id):
        """按ID查询航班"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM flights WHERE flight_id = %s", (flight_id,))
            return cursor.fetchone()
        except Error as e:
            st.error(f"查询航班时出错: {e}")
            return None

    def get_all_flights(self):
        """查询所有航班"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM flights")
            return cursor.fetchall()
        except Error as e:
            st.error(f"查询航班时出错: {e}")
            return []

    def update_flight(self, flight_id, departure_city=None, arrival_city=None,
                      departure_time=None, arrival_time=None, price=None,
                      seats_available=None, update_by=None):
        """更新航班信息"""
        try:
            cursor = self.connection.cursor()
            update_fields = []
            update_data = []

            if departure_city is not None:
                update_fields.append("departure_city = %s")
                update_data.append(departure_city)
            if arrival_city is not None:
                update_fields.append("arrival_city = %s")
                update_data.append(arrival_city)
            if departure_time is not None:
                update_fields.append("departure_time = %s")
                update_data.append(departure_time)
            if arrival_time is not None:
                update_fields.append("arrival_time = %s")
                update_data.append(arrival_time)
            if price is not None:
                update_fields.append("price = %s")
                update_data.append(price)
            if seats_available is not None:
                update_fields.append("seats_available = %s")
                update_data.append(seats_available)
            if update_by is not None:
                update_fields.append("update_by = %s")
                update_data.append(update_by)

            if not update_fields:
                print("没有提供更新字段")
                return False

            update_fields.append("update_time = CURRENT_TIMESTAMP")
            query = f"UPDATE flights SET {', '.join(update_fields)} WHERE flight_id = %s"
            update_data.append(flight_id)

            cursor.execute(query, tuple(update_data))
            self.connection.commit()
            print(f"航班ID {flight_id} 更新成功")
            return True
        except Error as e:
            print(f"更新航班时出错: {e}")
            self.connection.rollback()
            return False

    def delete_flight(self, flight_id):
        """删除航班"""
        try:
            cursor = self.connection.cursor()
            cursor.execute("DELETE FROM flights WHERE flight_id = %s", (flight_id,))
            self.connection.commit()
            return cursor.rowcount > 0
        except Error as e:
            st.error(f"删除航班时出错: {e}")
            return False

    # 预订相关操作
    def add_booking(self, user_id, flight_id, status, create_by=None):
        """添加新预订"""
        try:
            cursor = self.connection.cursor()
            query = """
            INSERT INTO bookings (user_id, flight_id, status, create_by)
            VALUES (%s, %s, %s, %s)
            """
            data = (user_id, flight_id, status, create_by)
            cursor.execute(query, data)
            self.connection.commit()
            print(f"预订添加成功，ID: {cursor.lastrowid}")
            return cursor.lastrowid
        except Error as e:
            print(f"添加预订时出错: {e}")
            return None

    def get_booking_by_id(self, booking_id):
        """按ID查询预订"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM bookings WHERE booking_id = %s", (booking_id,))
            return cursor.fetchone()
        except Error as e:
            st.error(f"查询预订时出错: {e}")
            return None

    def get_all_bookings(self):
        """查询所有预订"""
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM bookings")
            return cursor.fetchall()
        except Error as e:
            st.error(f"查询预订时出错: {e}")
            return []

    def update_booking(self, booking_id, status=None, update_by=None):
        """更新预订状态"""
        try:
            cursor = self.connection.cursor()
            update_fields = []
            update_data = []

            if status is not None:
                update_fields.append("status = %s")
                update_data.append(status)
            if update_by is not None:
                update_fields.append("update_by = %s")
                update_data.append(update_by)

            if not update_fields:
                print("没有提供更新字段")
                return False

            update_fields.append("update_time = CURRENT_TIMESTAMP")
            query = f"UPDATE bookings SET {', '.join(update_fields)} WHERE booking_id = %s"
            update_data.append(booking_id)

            cursor.execute(query, tuple(update_data))
            self.connection.commit()
            print(f"预订ID {booking_id} 更新成功")
            return True
        except Error as e:
            print(f"更新预订时出错: {e}")
            self.connection.rollback()
            return False

    def delete_booking(self, booking_id):
        """删除预订"""
        try:
            cursor = self.connection.cursor()
            cursor.execute("DELETE FROM bookings WHERE booking_id = %s", (booking_id,))
            self.connection.commit()
            return cursor.rowcount > 0
        except Error as e:
            st.error(f"删除预订时出错: {e}")
            return False


    def initialize_sample_data(self):
        """初始化示例数据"""
        try:
            cursor = self.connection.cursor(dictionary=True)

            # 检查是否已存在admin用户
            cursor.execute("SELECT user_id FROM users WHERE username = 'admin'")
            if cursor.fetchone():
                st.warning("测试数据已存在，无需重复初始化")
                return True

            # 添加管理员用户
            admin_id = self.add_user(
                username='admin',
                password='admin123',
                email='admin@example.com',
                first_name='Admin',
                last_name='User'
            )

            if admin_id:
                # 添加示例航班
                flight_id = self.add_flight(
                    departure_city='Beijing',
                    arrival_city='Shanghai',
                    departure_time=datetime(2023, 12, 1, 8, 0),
                    arrival_time=datetime(2023, 12, 1, 10, 30),
                    price=1200.00,
                    seats_available=150,
                    create_by=admin_id
                )

                if flight_id:
                    # 添加示例预订
                    booking_id = self.add_booking(
                        user_id=admin_id,
                        flight_id=flight_id,
                        status='confirmed',
                        create_by=admin_id
                    )
                    st.success("测试数据初始化成功！")
                    return True
            return False
        except Exception as e:
            st.error(f"初始化数据时出错: {e}")
            return False

# 初始化数据库连接
def init_connection():
    if 'db_system' not in st.session_state:
        db_config = {
            'host': "localhost",
            'database': "flight_system",
            'user': "root",
            'password': "JCX1105"
        }
        st.session_state.db_system = FlightBookingSystem(**db_config)
        if not st.session_state.db_system.connect():
            st.error("无法连接到数据库")
            st.stop()
    return st.session_state.db_system


def combine_datetime(date_part, time_part):
    """将日期和时间部分组合成datetime对象"""
    return datetime.combine(date_part, time_part)


# 用户管理页面
def user_management(system):
    st.header("📝 用户管理")
    tab1, tab2, tab3 = st.tabs(["查询用户", "添加用户", "管理用户"])

    with tab1:
        st.subheader("用户列表")
        users = system.get_all_users()
        if users:
            st.dataframe(pd.DataFrame(users))
        else:
            st.warning("没有用户数据")

    with tab2:
        st.subheader("添加新用户")
        with st.form("add_user_form"):
            col1, col2 = st.columns(2)
            with col1:
                username = st.text_input("用户名*", max_chars=50)
                email = st.text_input("邮箱")
                first_name = st.text_input("名字")
            with col2:
                password = st.text_input("密码*", type="password")
                phone = st.text_input("电话")
                last_name = st.text_input("姓氏")

            submitted = st.form_submit_button("添加用户")
            if submitted:
                if username and password:
                    result = system.add_user(
                        username=username,
                        password=password,
                        email=email,
                        phone=phone,
                        first_name=first_name,
                        last_name=last_name
                    )
                    if result:
                        st.success(f"用户添加成功！用户ID: {result}")
                        st.rerun()
                else:
                    st.error("用户名和密码是必填项")

    with tab3:
        st.subheader("管理用户")
        users = system.get_all_users()
        if users:
            user_options = {f"{u['user_id']} - {u['username']}": u['user_id'] for u in users}
            selected_user = st.selectbox("选择用户", options=list(user_options.keys()))
            user_id = user_options[selected_user]

            user = system.get_user_by_id(user_id)
            if user:
                with st.form(f"edit_user_{user_id}"):
                    col1, col2 = st.columns(2)
                    with col1:
                        new_username = st.text_input("用户名", value=user['username'])
                        new_email = st.text_input("邮箱", value=user.get('email', ''))
                        new_first_name = st.text_input("名字", value=user.get('first_name', ''))
                    with col2:
                        new_password = st.text_input("新密码 (留空不修改)", type="password")
                        new_phone = st.text_input("电话", value=user.get('phone', ''))
                        new_last_name = st.text_input("姓氏", value=user.get('last_name', ''))

                    submitted = st.form_submit_button("更新用户")
                    if submitted:
                        update_data = {
                            'username': new_username,
                            'email': new_email,
                            'phone': new_phone,
                            'first_name': new_first_name,
                            'last_name': new_last_name
                        }
                        if new_password:
                            update_data['password'] = new_password

                        if system.update_user(user_id, **update_data):
                            st.success("用户信息更新成功！")
                            st.rerun()

                # 删除按钮放在表单外部
                if st.button("删除用户", key=f"delete_user_{user_id}"):
                    if system.delete_user(user_id):
                        st.success("用户删除成功！")
                        st.rerun()


# 航班管理页面
def flight_management(system):
    st.header("✈️ 航班管理")
    tab1, tab2, tab3 = st.tabs(["查询航班", "添加航班", "管理航班"])

    with tab1:
        st.subheader("航班列表")
        flights = system.get_all_flights()
        if flights:
            st.dataframe(pd.DataFrame(flights))
        else:
            st.warning("没有航班数据")

    with tab2:
        st.subheader("添加新航班")
        with st.form("add_flight_form"):
            col1, col2 = st.columns(2)
            with col1:
                flight_number = st.text_input("航班号*")  # 新增字段
                departure_city = st.text_input("出发城市*")
                departure_date = st.date_input("出发日期*", min_value=date.today())
                departure_time = st.time_input("出发时间*", value=time(8, 0))
            with col2:
                arrival_city = st.text_input("到达城市*")
                arrival_date = st.date_input("到达日期*", min_value=date.today())
                arrival_time = st.time_input("到达时间*", value=time(10, 30))
                price = st.number_input("价格*", min_value=0.0, format="%.2f")
                seats_available = st.number_input("可用座位数*", min_value=1)

            submitted = st.form_submit_button("添加航班")
            if submitted:
                if flight_number and departure_city and arrival_city:  # 验证航班号
                    departure_datetime = combine_datetime(departure_date, departure_time)
                    arrival_datetime = combine_datetime(arrival_date, arrival_time)
                    result = system.add_flight(
                        flight_number=flight_number,  # 新增参数
                        departure_city=departure_city,
                        arrival_city=arrival_city,
                        departure_time=departure_datetime,
                        arrival_time=arrival_datetime,
                        price=price,
                        seats_available=seats_available
                    )
                    if result:
                        st.success(f"航班添加成功！航班ID: {result}")
                        st.rerun()
                else:
                    st.error("带*的字段是必填项")

    with tab3:
        st.subheader("管理航班")
        flights = system.get_all_flights()
        if flights:
            flight_options = {f"{f['flight_id']} - {f['departure_city']} to {f['arrival_city']}": f['flight_id'] for f
                              in flights}
            selected_flight = st.selectbox("选择航班", options=list(flight_options.keys()))
            flight_id = flight_options[selected_flight]

            flight = system.get_flight_by_id(flight_id)
            if flight:
                with st.form(f"edit_flight_{flight_id}"):
                    col1, col2 = st.columns(2)
                    with col1:
                        new_departure = st.text_input("出发城市", value=flight['departure_city'])
                        new_departure_dt = datetime.strptime(str(flight['departure_time']), '%Y-%m-%d %H:%M:%S')
                        new_departure_date = st.date_input("出发日期", value=new_departure_dt.date())
                        new_departure_time = st.time_input("出发时间", value=new_departure_dt.time())
                        new_price = st.number_input("价格", min_value=0.0, value=float(flight['price']), format="%.2f")
                    with col2:
                        new_arrival = st.text_input("到达城市", value=flight['arrival_city'])
                        new_arrival_dt = datetime.strptime(str(flight['arrival_time']), '%Y-%m-%d %H:%M:%S')
                        new_arrival_date = st.date_input("到达日期", value=new_arrival_dt.date())
                        new_arrival_time = st.time_input("到达时间", value=new_arrival_dt.time())
                        new_seats = st.number_input("可用座位数", min_value=1, value=flight['seats_available'])

                    submitted = st.form_submit_button("更新航班")
                    if submitted:
                        new_departure_datetime = combine_datetime(new_departure_date, new_departure_time)
                        new_arrival_datetime = combine_datetime(new_arrival_date, new_arrival_time)
                        if system.update_flight(
                                flight_id=flight_id,
                                departure_city=new_departure,
                                arrival_city=new_arrival,
                                departure_time=new_departure_datetime,
                                arrival_time=new_arrival_datetime,
                                price=new_price,
                                seats_available=new_seats
                        ):
                            st.success("航班信息更新成功！")
                            st.rerun()

                # 删除按钮放在表单外部
                if st.button("删除航班", key=f"delete_flight_{flight_id}"):
                    if system.delete_flight(flight_id):
                        st.success("航班删除成功！")
                        st.rerun()


# 预订管理页面
def booking_management(system):
    st.header("📅 预订管理")
    tab1, tab2, tab3 = st.tabs(["查询预订", "添加预订", "管理预订"])

    with tab1:
        st.subheader("预订列表")
        bookings = system.get_all_bookings()
        if bookings:
            bookings_df = pd.DataFrame(bookings)
            st.dataframe(bookings_df)
        else:
            st.warning("没有预订数据")

    with tab2:
        st.subheader("添加新预订")
        users = system.get_all_users()
        flights = system.get_all_flights()

        if users and flights:
            user_options = {f"{u['user_id']} - {u['username']}": u['user_id'] for u in users}
            flight_options = {f"{f['flight_id']} - {f['departure_city']} to {f['arrival_city']}": f['flight_id'] for f
                              in flights}

            with st.form("add_booking_form"):
                selected_user = st.selectbox("选择用户", options=list(user_options.keys()))
                selected_flight = st.selectbox("选择航班", options=list(flight_options.keys()))
                status = st.selectbox("状态", ["confirmed", "pending", "cancelled"])

                submitted = st.form_submit_button("添加预订")
                if submitted:
                    user_id = user_options[selected_user]
                    flight_id = flight_options[selected_flight]
                    result = system.add_booking(
                        user_id=user_id,
                        flight_id=flight_id,
                        status=status
                    )
                    if result:
                        st.success(f"预订添加成功！预订ID: {result}")
                        st.rerun()
        else:
            st.warning("请先添加用户和航班")

    with tab3:
        st.subheader("管理预订")
        bookings = system.get_all_bookings()
        if bookings:
            booking_options = {f"{b['booking_id']} - 用户ID: {b['user_id']}, 航班ID: {b['flight_id']}": b['booking_id']
                               for b in bookings}
            selected_booking = st.selectbox("选择预订", options=list(booking_options.keys()))
            booking_id = booking_options[selected_booking]

            booking = system.get_booking_by_id(booking_id)
            if booking:
                col1, col2 = st.columns(2)
                with col1:
                    st.write(f"**用户ID:** {booking['user_id']}")
                    st.write(f"**航班ID:** {booking['flight_id']}")
                with col2:
                    st.write(f"**状态:** {booking['status']}")
                    st.write(f"**创建时间:** {booking['create_time']}")

                # 统一状态选项
                status_options = ["confirmed", "pending", "cancelled"]
                current_status = booking['status'] if booking['status'] in status_options else "confirmed"
                new_status = st.selectbox("更新状态", status_options,
                                          index=status_options.index(current_status))

                if st.button("更新状态", key=f"update_status_{booking_id}"):
                    if system.update_booking(
                            booking_id=booking_id,
                            status=new_status
                    ):
                        st.success("预订状态更新成功！")
                        st.rerun()

                if st.button("删除预订", key=f"delete_booking_{booking_id}"):
                    if system.delete_booking(booking_id):
                        st.success("预订删除成功！")
                        st.rerun()


# 主应用
def main():
    st.set_page_config(page_title="航班预订系统", layout="wide", page_icon="✈️")
    st.title("✈️ 航班预订管理系统")

    # 初始化数据库连接
    system = init_connection()

    # 初始化数据按钮
    if st.sidebar.button("初始化实验数据"):
        if system.initialize_sample_data():
            st.sidebar.success("实验数据初始化成功！")
            st.rerun()
        else:
            st.sidebar.error("初始化数据失败")

    # 侧边栏导航
    menu = st.sidebar.radio("导航菜单",
                            ["用户管理", "航班管理", "预订管理"])

    if menu == "用户管理":
        user_management(system)
    elif menu == "航班管理":
        flight_management(system)
    elif menu == "预订管理":
        booking_management(system)

    # 在侧边栏显示数据库状态
    st.sidebar.markdown("---")
    if system.connection and system.connection.is_connected():
        st.sidebar.success("数据库已连接 ✅")
    else:
        st.sidebar.error("数据库未连接 ❌")


if __name__ == "__main__":
    main()